const MH_SHA256_H0:u32 = 0x6a09e667;
const MH_SHA256_H1:u32 = 0xbb67ae85;
const MH_SHA256_H2:u32 = 0x3c6ef372;
const MH_SHA256_H3:u32 = 0xa54ff53a;
const MH_SHA256_H4:u32 = 0x510e527f;
const MH_SHA256_H5:u32 = 0x9b05688c;
const MH_SHA256_H6:u32 = 0x1f83d9ab;
const MH_SHA256_H7:u32 = 0x5be0cd19;

const ISAL_SHA256_DIGEST_NWORDS:usize = 8;
const ISAL_SHA256_BLOCK_SIZE:usize = 64;

pub struct Sha256Ctx {
    total_length:usize,
    partial_block_buffer_length:usize,
    partial_block_buffer:[u8;ISAL_SHA256_BLOCK_SIZE * 2], 
    digest:[u32;ISAL_SHA256_DIGEST_NWORDS],
}

impl Sha256Ctx {
    pub fn new() -> Self {
        Self { 
            total_length:0, 
            partial_block_buffer_length:0,
            partial_block_buffer:[0;ISAL_SHA256_BLOCK_SIZE * 2], 
            digest:seq!(I in 0..8 { [#(MH_SHA256_H~I,)*] }),
        }
    }

    pub fn update(&mut self, data:&[u8]) {
        let mut input_data = data;
        // Advance byte counter
        self.total_length += data.len();

        // If there is anything currently buffered in the extra blocks, append to it until it
        // contains a whole block. Or if the user's buffer contains less than a whole block, append
        // as much as possible to the extra block.
        if (self.partial_block_buffer_length != 0) || (input_data.len() < ISAL_SHA256_BLOCK_SIZE) {
            // Compute how many bytes to copy from user buffer into extra block
            let copy_len = std::cmp::min(ISAL_SHA256_BLOCK_SIZE - self.partial_block_buffer_length, input_data.len());
            let partial_block_buffer_copy_end = self.partial_block_buffer_length + copy_len;

            if copy_len != 0 {
                self.partial_block_buffer[self.partial_block_buffer_length..partial_block_buffer_copy_end].copy_from_slice(&input_data[0..copy_len]);
                self.partial_block_buffer_length += copy_len;
                input_data = &input_data[copy_len..];
            }
            // The extra block should never contain more than 1 block here
            assert!(self.partial_block_buffer_length <= ISAL_SHA256_BLOCK_SIZE);

            // If the extra block buffer contains exactly 1 block, it can be hashed.
            if self.partial_block_buffer_length >= ISAL_SHA256_BLOCK_SIZE {
                    self.partial_block_buffer_length = 0;
                    Self::sha256_single(&self.partial_block_buffer, &mut self.digest);
            }
        }

        // If the extra blocks are empty, begin hashing what remains in the user's buffer.
        while input_data.len() >= ISAL_SHA256_BLOCK_SIZE {
            Self::sha256_single(input_data, &mut self.digest);
            input_data = &input_data[ISAL_SHA256_BLOCK_SIZE..];
        }

        if input_data.len() > 0 {
            let copy_len = input_data.len();
            self.partial_block_buffer[0..copy_len].copy_from_slice(input_data);
            self.partial_block_buffer_length = input_data.len();
        }
        return;
    }

    pub fn finalize(&mut self, digest:&mut [u8;std::mem::size_of::<u32>() * ISAL_SHA256_DIGEST_NWORDS]) {
        // Padding the first block
        self.partial_block_buffer[self.partial_block_buffer_length] = 0x80;
        self.partial_block_buffer_length += 1;

        self.partial_block_buffer[self.partial_block_buffer_length..ISAL_SHA256_BLOCK_SIZE].fill(0);

        // Calculate the first block without total_length if padding needs 2 block
        if self.partial_block_buffer_length > ISAL_SHA256_BLOCK_SIZE - 8 {
            Self::sha256_single(&mut self.partial_block_buffer, &mut self.digest);
            // Padding the second block
            self.partial_block_buffer[0..ISAL_SHA256_BLOCK_SIZE].fill(0);
        }

        // Padding the block
        let len_in_bit:[u8;std::mem::size_of::<u64>()] = (self.total_length as u64 * 8).to_be_bytes();
        self.partial_block_buffer[ISAL_SHA256_BLOCK_SIZE - 8..ISAL_SHA256_BLOCK_SIZE].copy_from_slice(&len_in_bit);
        Self::sha256_single(&self.partial_block_buffer, &mut self.digest);

        /* Output the digests of mh_sha256 */
        seq!(N in 0..8 { digest[N*std::mem::size_of::<u32>()..(N+1)*std::mem::size_of::<u32>()].copy_from_slice(&self.digest[N].to_be_bytes()); });
    }

    
    pub(crate) fn sha256_single(data:&[u8], digest:&mut [u32;ISAL_SHA256_DIGEST_NWORDS]) {
        let mut a = digest[0];
        let mut b = digest[1];
        let mut c = digest[2];
        let mut d = digest[3];
        let mut e = digest[4];
        let mut f = digest[5];
        let mut g = digest[6];
        let mut h = digest[7];

        let mut w = [0u32;16];
        macro_rules! W { ($i:expr) => { w[ $i & 0xf] }; }
        macro_rules! S0 { ($w:expr) => { $w.rotate_right(7) ^ $w.rotate_right(18) ^ ($w >> 3)};}
        macro_rules! S1 { ($w:expr) => { $w.rotate_right(17) ^ $w.rotate_right(19) ^ ($w >> 10)};}
        macro_rules! s0 { ($a:expr) => { $a.rotate_right(2) ^ $a.rotate_right(13) ^ $a.rotate_right(22) };}
        macro_rules! s1 { ($e:expr) => { $e.rotate_right(6) ^ $e.rotate_right(11) ^ $e.rotate_right(25) };}
        macro_rules! maj { ($a:expr, $b:expr, $c:expr) => { ($a & $b) ^ ($a & $c) ^ ($b & $c)};}
        macro_rules! ch { ($e:expr, $f:expr, $g:expr) => { ($e & $f) ^ ($g & !$e) };}
        macro_rules! step {
            ($i:expr, $a:expr, $b:expr, $c:expr, $d:expr, $e:expr, $f:expr, $g:expr, $h:expr, $k:expr) => {
                if $i < 16 {
                    W!($i) = u32::from_be_bytes(data[$i*std::mem::size_of::<u32>()..($i+1)*std::mem::size_of::<u32>()].try_into().unwrap());
                } else {
                    W!($i) = W!($i - 16).wrapping_add(S0!(W!($i - 15))).wrapping_add(W!($i - 7)).wrapping_add(S1!(W!($i - 2)));
                }
        
                let t2 = s0!($a).wrapping_add(maj!($a, $b, $c));
                let t1 = $h.wrapping_add(s1!($e)).wrapping_add(ch!($e, $f, $g)).wrapping_add($k).wrapping_add(W!($i));

                $d = $d.wrapping_add(t1);
                $h = t1.wrapping_add(t2);
            };
        }
        
        step!(0, a, b, c, d, e, f, g, h, 0x428a2f98);
        step!(1, h, a, b, c, d, e, f, g, 0x71374491);
        step!(2, g, h, a, b, c, d, e, f, 0xb5c0fbcf);
        step!(3, f, g, h, a, b, c, d, e, 0xe9b5dba5);
        step!(4, e, f, g, h, a, b, c, d, 0x3956c25b);
        step!(5, d, e, f, g, h, a, b, c, 0x59f111f1);
        step!(6, c, d, e, f, g, h, a, b, 0x923f82a4);
        step!(7, b, c, d, e, f, g, h, a, 0xab1c5ed5);
        step!(8, a, b, c, d, e, f, g, h, 0xd807aa98);
        step!(9, h, a, b, c, d, e, f, g, 0x12835b01);
        step!(10, g, h, a, b, c, d, e, f, 0x243185be);
        step!(11, f, g, h, a, b, c, d, e, 0x550c7dc3);
        step!(12, e, f, g, h, a, b, c, d, 0x72be5d74);
        step!(13, d, e, f, g, h, a, b, c, 0x80deb1fe);
        step!(14, c, d, e, f, g, h, a, b, 0x9bdc06a7);
        step!(15, b, c, d, e, f, g, h, a, 0xc19bf174);
        step!(16, a, b, c, d, e, f, g, h, 0xe49b69c1);
        step!(17, h, a, b, c, d, e, f, g, 0xefbe4786);
        step!(18, g, h, a, b, c, d, e, f, 0x0fc19dc6);
        step!(19, f, g, h, a, b, c, d, e, 0x240ca1cc);
        step!(20, e, f, g, h, a, b, c, d, 0x2de92c6f);
        step!(21, d, e, f, g, h, a, b, c, 0x4a7484aa);
        step!(22, c, d, e, f, g, h, a, b, 0x5cb0a9dc);
        step!(23, b, c, d, e, f, g, h, a, 0x76f988da);
        step!(24, a, b, c, d, e, f, g, h, 0x983e5152);
        step!(25, h, a, b, c, d, e, f, g, 0xa831c66d);
        step!(26, g, h, a, b, c, d, e, f, 0xb00327c8);
        step!(27, f, g, h, a, b, c, d, e, 0xbf597fc7);
        step!(28, e, f, g, h, a, b, c, d, 0xc6e00bf3);
        step!(29, d, e, f, g, h, a, b, c, 0xd5a79147);
        step!(30, c, d, e, f, g, h, a, b, 0x06ca6351);
        step!(31, b, c, d, e, f, g, h, a, 0x14292967);
        step!(32, a, b, c, d, e, f, g, h, 0x27b70a85);
        step!(33, h, a, b, c, d, e, f, g, 0x2e1b2138);
        step!(34, g, h, a, b, c, d, e, f, 0x4d2c6dfc);
        step!(35, f, g, h, a, b, c, d, e, 0x53380d13);
        step!(36, e, f, g, h, a, b, c, d, 0x650a7354);
        step!(37, d, e, f, g, h, a, b, c, 0x766a0abb);
        step!(38, c, d, e, f, g, h, a, b, 0x81c2c92e);
        step!(39, b, c, d, e, f, g, h, a, 0x92722c85);
        step!(40, a, b, c, d, e, f, g, h, 0xa2bfe8a1);
        step!(41, h, a, b, c, d, e, f, g, 0xa81a664b);
        step!(42, g, h, a, b, c, d, e, f, 0xc24b8b70);
        step!(43, f, g, h, a, b, c, d, e, 0xc76c51a3);
        step!(44, e, f, g, h, a, b, c, d, 0xd192e819);
        step!(45, d, e, f, g, h, a, b, c, 0xd6990624);
        step!(46, c, d, e, f, g, h, a, b, 0xf40e3585);
        step!(47, b, c, d, e, f, g, h, a, 0x106aa070);
        step!(48, a, b, c, d, e, f, g, h, 0x19a4c116);
        step!(49, h, a, b, c, d, e, f, g, 0x1e376c08);
        step!(50, g, h, a, b, c, d, e, f, 0x2748774c);
        step!(51, f, g, h, a, b, c, d, e, 0x34b0bcb5);
        step!(52, e, f, g, h, a, b, c, d, 0x391c0cb3);
        step!(53, d, e, f, g, h, a, b, c, 0x4ed8aa4a);
        step!(54, c, d, e, f, g, h, a, b, 0x5b9cca4f);
        step!(55, b, c, d, e, f, g, h, a, 0x682e6ff3);
        step!(56, a, b, c, d, e, f, g, h, 0x748f82ee);
        step!(57, h, a, b, c, d, e, f, g, 0x78a5636f);
        step!(58, g, h, a, b, c, d, e, f, 0x84c87814);
        step!(59, f, g, h, a, b, c, d, e, 0x8cc70208);
        step!(60, e, f, g, h, a, b, c, d, 0x90befffa);
        step!(61, d, e, f, g, h, a, b, c, 0xa4506ceb);
        step!(62, c, d, e, f, g, h, a, b, 0xbef9a3f7);
        step!(63, b, c, d, e, f, g, h, a, 0xc67178f2);

        digest[0] = digest[0].wrapping_add(a);
        digest[1] = digest[1].wrapping_add(b);
        digest[2] = digest[2].wrapping_add(c);
        digest[3] = digest[3].wrapping_add(d);
        digest[4] = digest[4].wrapping_add(e);
        digest[5] = digest[5].wrapping_add(f);
        digest[6] = digest[6].wrapping_add(g);
        digest[7] = digest[7].wrapping_add(h);
    }
}